package com.linyaonan.leetcode.medium._2583;


import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.Map;
import java.util.PriorityQueue;

/**
 * 给你一棵二叉树的根节点 root 和一个正整数 k 。
 *
 * 树中的 层和 是指 同一层 上节点值的总和。
 *
 * 返回树中第 k 大的层和（不一定不同）。如果树少于 k 层，则返回 -1 。
 *
 * 注意，如果两个节点与根节点的距离相同，则认为它们在同一层。
 *
 * 树中的节点数为 n
 * 2 <= n <= 10^5
 * 1 <= Node.val <= 10^6
 * 1 <= k <= n
 *
 * @author: Lin
 * @date: 2024/2/23
 */
public class KthLargestSumInABinaryTree {



    public long kthLargestLevelSum(TreeNode root, int k) {
        LinkedList<TreeNode> list = new LinkedList<>();
        // 每一层对应的数值大小
        HashMap<Integer, Long> map = new HashMap<>();
        HashMap<TreeNode, Integer> nodeWithCeng = new HashMap<>();

        list.addLast(root);
        nodeWithCeng.put(root, 0);

        while (!list.isEmpty()) {
            TreeNode t = list.removeFirst();
            Integer ceng = nodeWithCeng.get(t);
            map.put(ceng, map.getOrDefault(ceng, 0L) + t.val);

            if (t.left != null) {
                list.addLast(t.left);
                nodeWithCeng.put(t.left, ceng + 1);
            }
            if (t.right != null) {
                list.addLast(t.right);
                nodeWithCeng.put(t.right, ceng + 1);
            }
        }

        if (k > map.size()) {
            // 异常情况
            return -1L;
        }

        PriorityQueue<Long> queue = new PriorityQueue<>(map.size(), (o1, o2) -> Long.compare(o2, o1));

        for (Map.Entry<Integer, Long> entry : map.entrySet()) {
            queue.add(entry.getValue());
        }

        for (int i = 0; i < k - 1; i++) {
            queue.poll();
        }

        return queue.peek();
    }

    public long kthLargestLevelSum2(TreeNode root, int k) {

        LinkedList<TreeNode> list = new LinkedList<>();
        LinkedList<TreeNode> next = new LinkedList<>();
        list.addLast(root);

        ArrayList<Long> sums = new ArrayList<>();

        while (!list.isEmpty()) {
            long sum = 0;

            for (TreeNode treeNode : list) {
                sum += treeNode.val;
                if (treeNode.left != null) {
                    next.addLast(treeNode.left);
                }
                if (treeNode.right != null) {
                    next.addLast(treeNode.right);
                }
            }
            sums.add(sum);
            list = next;
            next.clear();
        }

        if (sums.size() < k) {
            return -1;
        }
        Collections.sort(sums);
        return sums.get(sums.size() - k);
    }


}
